#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
健康顾问角色扮演任务Agent
使用camel-ai框架实现一个健康顾问，为患者提供饮食和锻炼建议
"""

from camel.societies import RolePlaying
from camel.models import ModelFactory
from camel.types import ModelPlatformType
from camel.utils import print_text_animated
from colorama import Fore, Style
import os
import re
import time
from dotenv import load_dotenv

load_dotenv()
api_key = os.getenv('MODELSCOPE_SDK_TOKEN')

def create_health_advisor_model():
    """创建健康顾问使用的模型"""
    model = ModelFactory.create(
        model_platform=ModelPlatformType.OPENAI_COMPATIBLE_MODEL,
        model_type="Qwen/Qwen2.5-72B-Instruct",
        url='https://api-inference.modelscope.cn/v1/',
        api_key=api_key
    )
    return model

def setup_role_playing_session():
    """设置角色扮演会话"""
    # 创建模型
    model = create_health_advisor_model()
    
    # 定义任务提示
    task_prompt = (
        "为一位想要改善健康状况的患者提供全面的饮食和锻炼建议。"
        "请使用思维链方式分解整个建议过程，逐步分析患者需求，"
        "然后提供详细的饮食和锻炼方案。"
        "最终方案以markdown格式输出。"
    )
    
    # 设置角色扮演
    role_play_session = RolePlaying(
        assistant_role_name="健康顾问",
        user_role_name="患者",
        assistant_agent_kwargs=dict(model=model),
        user_agent_kwargs=dict(model=model),
        task_prompt=task_prompt,
        with_task_specify=True,
        task_specify_agent_kwargs=dict(model=model),
        output_language='中文'
    )
    
    return role_play_session

def print_session_info(role_play_session, task_prompt):
    """打印会话信息"""
    print(
        Fore.GREEN + Style.BRIGHT
        + f"健康顾问系统消息:\n{role_play_session.assistant_sys_msg}\n"
    )
    print(
        Fore.BLUE + Style.BRIGHT
        + f"患者系统消息:\n{role_play_session.user_sys_msg}\n"
    )

    print(Fore.YELLOW + f"原始任务提示:\n{task_prompt}\n")
    print(
        Fore.CYAN + Style.BRIGHT
        + "指定的任务提示:"
        + f"\n{role_play_session.specified_task_prompt}\n"
    )
    print(Fore.RED + f"最终任务提示:\n{role_play_session.task_prompt}\n")

def save_markdown_to_file(markdown_content):
    """将markdown内容保存到文件"""
    # 获取当前Unix时间戳
    timestamp = int(time.time())
    filename = f"health_advisor_{timestamp}.md"
    
    # 保存文件
    with open(filename, 'w', encoding='utf-8') as f:
        f.write(markdown_content)
    
    print(Fore.YELLOW + f"\n健康建议已保存到文件: {filename}\n" + Style.RESET_ALL)
    return filename

def extract_markdown_content(conversation_history):
    """从对话历史中提取markdown内容"""
    # 合并所有健康顾问的回复
    full_response = "\n".join(conversation_history)
    
    # 查找包含markdown格式的内容（查找以#开头的标题）
    # 这通常表示方案的开始部分
    pattern = r'(#\s+.*?)(?=Instruction:|$)'
    match = re.search(pattern, full_response, re.DOTALL)
    
    if match:
        # 如果找到markdown内容，提取它
        markdown_content = match.group(1).strip()
        return markdown_content
    
    # 如果没有找到特定格式，返回整个健康顾问的回复
    return full_response

def run_health_consultation(role_play_session, chat_turn_limit=15):
    """运行健康咨询对话"""
    print(Fore.MAGENTA + "="*60)
    print(Fore.MAGENTA + "开始健康咨询...")
    print(Fore.MAGENTA + "="*60 + Style.RESET_ALL)
    
    n = 0
    input_msg = role_play_session.init_chat()
    assistant_responses = []  # 存储健康顾问的回复
    
    while n < chat_turn_limit:
        n += 1
        assistant_response, user_response = role_play_session.step(input_msg)

        if assistant_response.terminated:
            print(
                Fore.GREEN
                + (
                    "健康顾问终止咨询。原因: "
                    f"{assistant_response.info['termination_reasons']}."
                )
            )
            break
        if user_response.terminated:
            print(
                Fore.GREEN
                + (
                    "患者结束咨询。"
                    f"原因: {user_response.info['termination_reasons']}."
                )
            )
            break

        # 打印患者的输出
        print_text_animated(
            Fore.BLUE + Style.BRIGHT + f"患者:\n\n{user_response.msg.content}\n"
        )

        if "CAMEL_TASK_DONE" in user_response.msg.content:
            # 提取并保存markdown内容
            markdown_content = extract_markdown_content(assistant_responses)
            if markdown_content:
                save_markdown_to_file(markdown_content)
            
            print(Fore.MAGENTA + "\n" + "="*60)
            print(Fore.MAGENTA + "健康咨询完成!")
            print(Fore.MAGENTA + "="*60 + Style.RESET_ALL)
            break

        # 打印健康顾问的输出
        print_text_animated(
            Fore.GREEN + Style.BRIGHT + "健康顾问:\n\n"
            f"{assistant_response.msg.content}\n"
        )
        
        # 保存健康顾问的回复
        assistant_responses.append(assistant_response.msg.content)

        input_msg = assistant_response.msg

def main():
    """主函数"""
    print(Fore.MAGENTA + Style.BRIGHT + "健康顾问角色扮演任务Agent")
    print(Fore.MAGENTA + "="*60 + Style.RESET_ALL)
    
    # 设置角色扮演会话
    role_play_session = setup_role_playing_session()
    
    # 获取任务提示
    task_prompt = (
        "为一位想要改善健康状况的患者提供全面的饮食和锻炼建议。"
        "请使用思维链方式分解整个建议过程，逐步分析患者需求，"
        "然后提供详细的饮食和锻炼方案。"
        "最终方案以markdown格式输出。"
    )
    
    # 打印会话信息
    print_session_info(role_play_session, task_prompt)
    
    # 运行健康咨询对话
    run_health_consultation(role_play_session)

if __name__ == "__main__":
    main()